import os
import json

import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from vit_model import vit_base_patch16_224_in21k as create_model
from tqdm import tqdm
from my_dataset import MyDataSet
from sklearn.metrics import f1_score
import numpy as np


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose(
        [transforms.RandomResizedCrop(224),
         transforms.RandomHorizontalFlip(),
         transforms.ToTensor(),
         transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

    val_slides = {'0': ['02', '07', '13', '16', '17', '18', '27', '35', '37', '46'],
                  '1': ['20', '24', '01', '05', '06', '12', '30', '32', '34', '42', '44', '47']}
    slide_pre_acc_num = 0
    for label in val_slides.keys():
        print(f"测试{'MSI-H' if label == '1' else 'MSS'}样本：")
        for slide in val_slides[label]:
            source_path = os.path.join('./datasets', label, slide)
            test_images_path = []
            test_num = 0  # 参与测试的样本数
            for img_file in os.listdir(source_path):
                test_images_path.append(os.path.join(source_path, img_file))
                test_num += 1
            test_images_label = [int(label)] * test_num
            # 实例化验证数据集
            test_dataset = MyDataSet(images_path=test_images_path,
                                    images_class=test_images_label,
                                    transform=data_transform)

            test_loader = torch.utils.data.DataLoader(test_dataset,
                                                       batch_size=8,
                                                       shuffle=False,
                                                       pin_memory=True,
                                                       num_workers=8,
                                                       collate_fn=test_dataset.collate_fn)

            accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数
            # create model
            model = create_model(num_classes=2, has_logits=False).to(device)
            # load model weights
            model_weight_path = "./weights/model-3.pth"
            model.load_state_dict(torch.load(model_weight_path, map_location=device))
            model.eval()

            sample_num = 0
            conf = 0  # 设置置信度，只取高置信度的图块预测结果
            preds = torch.tensor([]).to(device)
            labels_all = torch.tensor([]).to(device)
            msi_h_pre_num = 0
            for data in tqdm(test_loader):
                images, labels = data

                pred = model(images.to(device))
                pred = torch.max(F.softmax(pred, dim=1), dim=1)
                conf_mask = pred[0].gt(conf)
                pred_classes = pred[1][conf_mask]
                labels = labels.to(device)[conf_mask]
                accu_num += torch.eq(pred_classes, labels).sum()
                sample_num += labels.shape[0]

                preds = torch.cat((preds, pred_classes))
                labels_all = torch.cat((labels_all, labels))

                msi_h_pre_num += np.count_nonzero(pred_classes.to('cpu'))

            print(f"测试切片{slide}:")
            print(f'测试acc：{accu_num/sample_num}')

            # 二分类任务f1_score只输出一个类的计算结果，这个类由pos_label决定
            print(f"测试F1-score：{f1_score(labels_all.to('cpu'), preds.to('cpu'), pos_label=int(label))}")

            # 设置一个比例rate，当一个切片中预测为MSI-H的图块数量超过这个比例时即判定为MSI-H，否则为MSS
            rate = 0.25
            slide_pre = '1' if msi_h_pre_num > len(test_loader) * 8 * rate else '0'
            slide_pre_acc_num += 1 if slide_pre == label else 0
            print(f"该切片预测为{'MSI-H' if slide_pre == '1' else 'MSS'}")

    slide_num = len(val_slides['0']) + len(val_slides['1'])
    print()
    print(f"切片预测准确率acc：{slide_pre_acc_num / slide_num}")


if __name__ == '__main__':
    main()
